# import
import numpy as np
import pandas as pd
from sklearn.model_selection import KFold
from sklearn.utils import resample
from numba import njit, jit

import time
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from util.models import *
from util.misc import *
from completeness import *

## mat file
from scipy.io import loadmat

if __name__ == '__main__':
    # seed
    seed = 20220219
    np.random.seed(seed)
    print('seed={}'.format(seed))

    # load dataset
    # dir_name = os.path.dirname(__file__)
    # matdata = loadmat(dir_name + '/Bruhin_et_al_2010.mat')
    base = os.path.dirname(os.path.abspath(__file__))
    file_name = os.path.normpath(os.path.join(base, '../data/Bruhin_et_al_2010.mat'))
    matdata = loadmat(file_name)

    df_full = pd.DataFrame()
    for key in matdata.keys():
        if not '__' in key:
            df_full[key] = matdata[key].flatten()

    ## drop lotteries over losses
    df = df_full[df_full['z2'] >= 0]

    # preprocessing
    X, Y, subject_idx, cov_idx, cov_list = preprocess(df)


    # main
    start_time = time.perf_counter()

    ## alpha, delta, gamma
    grid_size_alpha = 0.01
    grid_size_gamma = 0.01
    grid_size_delta = 0.1
    alpha_grid = np.arange(grid_size_alpha, 1+grid_size_alpha, grid_size_alpha)
    gamma_grid = np.arange(grid_size_gamma, 1+grid_size_gamma, grid_size_gamma)
    delta_grid = np.arange(grid_size_delta, 5+grid_size_delta, grid_size_delta)
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)

    print('CPT(alpha, delta, gamma)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)
    # se_bootstrap = bootstrap_se(X, Y, subject_idx, cov_idx, cov_list, param_grid, bs_sample_size=5)

    # print
    # print('stderr (bootstrapped): {}'.format(se_bootstrap))

    ## alpha, delta
    alpha_grid = np.arange(grid_size_alpha, 1+grid_size_alpha, grid_size_alpha)
    gamma_grid = np.array([1.0])
    delta_grid = np.arange(grid_size_delta, 5+grid_size_delta, grid_size_delta)
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)
    print('CPT(alpha, delta)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)

    ## alpha, gamma
    alpha_grid = np.arange(grid_size_alpha, 1+grid_size_alpha, grid_size_alpha)
    gamma_grid = np.arange(grid_size_gamma, 1+grid_size_gamma, grid_size_gamma)
    delta_grid = np.array([1.0])
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)
    print('CPT(alpha, gamma)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)

    ## delta, gamma
    alpha_grid = np.array([1.0])
    gamma_grid = np.arange(grid_size_gamma, 1+grid_size_gamma, grid_size_gamma)
    delta_grid = np.arange(grid_size_delta, 5+grid_size_delta, grid_size_delta)
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)
    print('CPT(delta, gamma)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)

    ## alpha
    alpha_grid = np.arange(grid_size_alpha, 1+grid_size_alpha, grid_size_alpha)
    gamma_grid = np.array([1.0])
    delta_grid = np.array([1.0])
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)
    print('CPT(alpha)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)

    ## delta
    alpha_grid = np.array([1.0])
    gamma_grid = np.array([1.0])
    delta_grid = np.arange(grid_size_delta, 5+grid_size_delta, grid_size_delta)
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)
    print('CPT(delta)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)

    ## gamma
    alpha_grid = np.array([1.0])
    gamma_grid = np.arange(grid_size_gamma, 1+grid_size_gamma, grid_size_gamma)
    delta_grid = np.array([1.0])
    param_grid = cartesian_product(alpha_grid, gamma_grid, delta_grid)
    print('CPT(gamma)')
    completeness, stderr, model_best_params =\
        cross_validation(X, Y, subject_idx, cov_idx, cov_list, param_grid)

    # time
    end_time = time.perf_counter()
    elapsed_time = end_time - start_time
    hour = elapsed_time//3600
    elapsed_time = elapsed_time - hour*3600
    minute = elapsed_time//60
    elapsed_time = elapsed_time - minute*60
    second = elapsed_time

    print('')
    print('process time: {}:{}:{}'.format(int(hour), int(minute), int(second)))
